#include <stdio.h>
#include <GL/glew.h>
#include <Camera/FlyCamera.h>
#include <Core/Assert.h>
#include <Core/Timer.h>
#include <Graphics/Device.h>
#include <Graphics/Shader.h>
#include <Graphics/VertexBuffer.h>
#include <Maths/Constants.h>
#include <Maths/MinMax.h>
#include <Maths/Projections.h>
#include <Input/Keyboard.h>
#include <Input/Mouse.h>
#include <vector>
#include <hash_map>

#include "intersections.h"
#include "structs.h"


using namespace std;
using namespace stdext;


const size_t kScreenWidth = 1280;
const size_t kScreenHeight = 720;
const float kAspectRatio = (float)kScreenWidth / (float)kScreenHeight;


enum Mode
{
    kClusterColour,
    kClusterFormFactor,
    kClusterIrradiance,
    kClusterContribution,
    kModeCount
};


const float kPatchSize = 0.3f;
const size_t kInvalidIndex = 0xffffffff;

VertexBuffer* vb = NULL;
Shader* shader = NULL;

bool isPaused = false;
bool isBruteForce = false;
bool isDone = false;
bool isIndirectLightingEnabled = true;

float bounceAmount = 1;

double elapsedTime = 0;

FlyCamera camera(Vector3(0, 4.0f, 6.0f), Vector3::kZero);

vector<Patch*> allPatches;
vector<ContributionDebug>* contributionDebug;
size_t leafPatchCount = 0;
size_t rootPatchIdx = 0;
size_t selectedPatchIdx = kInvalidIndex;

Mode debugMode = kClusterColour;

vector<Quad> faces;
vector<Quad> patchQuads;

Vector3 lightPos;

vector<Vertex> vertices;

Matrix44 invProj = Matrix44::kIdentity;
Matrix44 invView = Matrix44::kIdentity;
Matrix44 invViewProj = Matrix44::kIdentity;
Matrix44 proj = Matrix44::kIdentity;
Matrix44 view = Matrix44::kIdentity;
Matrix44 viewProj = Matrix44::kIdentity;


Vector3 Bilerp(Vector3 a, Vector3 b, Vector3 c, Vector3 d, float u, float v)
{
    Vector3 ab = a + u * (b - a);
    Vector3 dc = d + u * (c - d);
    return ab + v * (dc - ab);
}


void AddQuad(Vector3Arg colour, const Quad& quad)
{
    // Flip winding order to raycast against back faces
    faces.push_back(Quad(quad.d, quad.c, quad.b, quad.a));

    Vector3 a = quad.a;
    Vector3 b = quad.b;
    Vector3 c = quad.c;
    Vector3 d = quad.d;

    float uLength = 0.5f * (Distance(a, b) + Distance(c, d));
    float vLength = 0.5f * (Distance(b, c) + Distance(d, a));

    int uSubdivisions = Max((int)(uLength / kPatchSize + 0.5f), 1);
    int vSubdivisions = Max((int)(vLength / kPatchSize + 0.5f), 1);

    for (int j = 0; j < vSubdivisions; ++j)
    {
        float v0 = (float)j / (float)vSubdivisions;
        float v1 = (float)(j+1) / (float)vSubdivisions;
        for (int i = 0; i < uSubdivisions; ++i)
        {
            float u0 = (float)i / (float)uSubdivisions;
            float u1 = (float)(i+1) / (float)uSubdivisions;

            Vector3 aa = Bilerp(a, b, c, d, u0, v0);
            Vector3 bb = Bilerp(a, b, c, d, u0, v1);
            Vector3 cc = Bilerp(a, b, c, d, u1, v1);
            Vector3 dd = Bilerp(a, b, c, d, u1, v0);

            vertices.push_back(Vertex(aa, colour));
            vertices.push_back(Vertex(cc, colour));
            vertices.push_back(Vertex(bb, colour));

            vertices.push_back(Vertex(aa, colour));
            vertices.push_back(Vertex(dd, colour));
            vertices.push_back(Vertex(cc, colour));

            Vector3 p = 0.25f * (aa + bb + cc + dd);
            Vector3 ab = bb - aa;
            Vector3 bc = cc - bb;
            Vector3 cd = dd - cc;
            Vector3 n = Normalize(Cross(b - a, c - a));
            float area = 0.5f * Dot(n, Cross(bc, ab) + Cross(cd, bc));
            allPatches.push_back(new Patch(colour, p, n, area));
            patchQuads.push_back(Quad(bb, aa, dd, cc));
        }
    }
}


float CalculateClusterScore(const Patch& cluster, const Patch& patch)
{
    float distanceScore = Distance(patch.position, cluster.position);
    float normalScore = 2.0f - Dot(patch.normal, cluster.normal);
    return normalScore * distanceScore;
}


void ClusterPatches(size_t patchBegin, size_t patchEnd, size_t clusterBegin, size_t clusterEnd)
{
    // Clear cluster assignments
    for (size_t i=clusterBegin; i!=clusterEnd; ++i)
        allPatches[i]->children.clear();

    // Assign patches to best cluster
    for (size_t patchIdx=patchBegin; patchIdx!=patchEnd; ++patchIdx)
    {
        Patch* patch = allPatches[patchIdx];

        float bestScore = kMaxPosFloat;
        Patch* bestCluster = allPatches[clusterBegin];
        for (size_t clusterIdx=clusterBegin; clusterIdx!=clusterEnd; ++clusterIdx)
        {
            Patch* cluster = allPatches[clusterIdx];
            float score = CalculateClusterScore(*cluster, *patch);
            if (score < bestScore)
            {
                bestCluster = cluster;
                bestScore = score;
            }
        }

        bestCluster->children.push_back(patchIdx);
    }

    // Recalculate cluster positions & normals, set children's parent reference
    for (size_t clusterIdx=clusterBegin; clusterIdx!=clusterEnd; ++clusterIdx)
    {
        Patch* cluster = allPatches[clusterIdx];

        Vector3 position = Vector3::kZero;
        Vector3 normal = Vector3::kZero;
        Vector3 albedo = Vector3::kZero;
        float area = 0;
        for (vector<size_t>::iterator ip=cluster->children.begin(), ep=cluster->children.end(); ip!=ep; ++ip)
        {
            Patch* patch = allPatches[*ip];
            area += patch->area;
            position += patch->position * patch->area;
            normal += patch->normal * patch->area;
            albedo += patch->albedo * patch->area;
            patch->parent = cluster;
        }

        if (area > 0)
        {
            cluster->position = position / area;
            cluster->normal = Normalize(normal);
            cluster->albedo = albedo / area;
        }
        else
        {
            cluster->position = position;
            cluster->normal = normal;
            cluster->albedo = albedo;
        }
        cluster->area = area;
    }
}


void ClusterPatches()
{
    printf("Clustering patches...\n");

    size_t beginIdx = 0;
    size_t endIdx = allPatches.size();
    do
    {
        // Seed k points for clustering
        size_t step = 2;
        size_t count = endIdx - beginIdx;
        for (size_t i = beginIdx; i < endIdx; i += step)
        {
            size_t idx;
            do
            {
                idx = i + (rand() % step);
            } while (idx >= endIdx);
            Patch* patch = allPatches[idx];
            allPatches.push_back(new Patch(patch->position, patch->normal, patch->area));
        }

        // k-means clustering (fixed number of iterations because I'm lazy)
        for (size_t pass = 0; pass < 8; ++pass)
            ClusterPatches(beginIdx, endIdx, endIdx, allPatches.size());

        beginIdx = endIdx;
        endIdx = allPatches.size();
    } while (endIdx-beginIdx > 1);

    rootPatchIdx = allPatches.size() - 1;
}


void CalculateContributingPatchesRecursive(Patch& leafPatch, size_t patchIdx, float* contributingPatches)
{
    const Patch& patch = *allPatches[patchIdx];
    if (&patch == &leafPatch)
        return;

    for (vector<size_t>::const_iterator i=patch.children.begin(), e=patch.children.end(); i!=e; ++i)
        CalculateContributingPatchesRecursive(leafPatch, *i, contributingPatches);

    bool areAnyChildrenClose = false;
    bool areAnyChildrenContributing = false;
    for (vector<size_t>::const_iterator i=patch.children.begin(), e=patch.children.end(); i!=e; ++i)
    {
        size_t childIdx = *i;
        const Patch* child = allPatches[childIdx];
        float contribution = contributingPatches[childIdx];
        if (contribution != 0)
        {
            areAnyChildrenContributing = true;
            float distanceSqr = DistanceSqr(leafPatch.position, child->position);
            if (distanceSqr <= 9 * patch.area)
            {
                areAnyChildrenClose = true;
                break;
            }
        }
    }

    if (areAnyChildrenContributing)
    {
        if (areAnyChildrenClose)
        {
            for (vector<size_t>::const_iterator i=patch.children.begin(), e=patch.children.end(); i!=e; ++i)
            {
                size_t childIdx = *i;
                float contribution = contributingPatches[childIdx];
                if (contribution != 0)
                    leafPatch.contributions.push_back(Contribution(childIdx, contribution));
            }
        }
        else
        {
            float totalContribution = 0;
            for (vector<size_t>::const_iterator i=patch.children.begin(), e=patch.children.end(); i!=e; ++i)
                totalContribution += contributingPatches[*i];
            contributingPatches[patchIdx] = totalContribution;
        }
    }
}


void CalculateContributingPatches()
{
    printf("Calculating contributing patches...\n");
    size_t totalPatchCount = allPatches.size();

    float* contributingPatches = new float[totalPatchCount];
    for (size_t i = 0; i < leafPatchCount; ++i)
    {
        if ((leafPatchCount-i) % 250 == 0)
            printf("   %d to go...\n", leafPatchCount-i);

        memset(contributingPatches, 0, totalPatchCount*sizeof(float));
        Patch& p = *allPatches[i];

        // Calculate form factor for each leaf patch
        for (size_t j = 0; j < leafPatchCount; ++j)
        {
            const Patch& q = *allPatches[j];
            if (&p == &q)
                continue;

            Vector3 pq = Normalize(q.position - p.position);
            float pn = Dot(pq, p.normal);
            if (pn <= 0)
                continue;

            float qn = -Dot(pq, q.normal);
            if (qn <= 0)
                continue;

            if (!IntersectRayQuads(p.position, q.position, &faces[0], faces.size()))
                contributingPatches[j] = pn * qn * q.area / (kPi * DistanceSqr(p.position, q.position));
        }

        CalculateContributingPatchesRecursive(p, rootPatchIdx, contributingPatches);
    }
    delete[] contributingPatches;

    size_t totalContributingPatches = 0;
    for (size_t i = 0; i < totalPatchCount; ++i)
    {
        Patch& p = *allPatches[i];
        totalContributingPatches += p.contributions.size();
    }
    printf("Average contributing patches = %.2f\n", (float)totalContributingPatches/(float)totalPatchCount);
}


void CalculateContributionDebug(size_t patchIdx, vector<ContributionDebug>& contributionDebug, const Contribution& contribution)
{
    const Patch& patch = *allPatches[patchIdx];
    if (patch.children.empty())
    {
        contributionDebug.push_back(ContributionDebug(patchIdx, contribution));
    }
    else
    {
        for (vector<size_t>::const_iterator ic=patch.children.begin(), ec=patch.children.end(); ic!=ec; ++ic)
            CalculateContributionDebug(*ic, contributionDebug, contribution);
    }
}


void InitScene()
{
    shader = Shader::Compile("Shader.glsl");
    Assert(shader->IsValid(), shader->GetError());

    Vector3 white = Vector3(0.75f, 0.75f, 0.75f);
    Vector3 red = Vector3(0.75f, 0.1f, 0.05f);
    Vector3 green = Vector3(0.05f, 0.75f, 0.1f);
    Vector3 blue = Vector3(0.05f, 0.1f, 0.75f);

    AddQuad(white,  Quad(Vector3(-5,  0, -5), Vector3(-5,  0,  5), Vector3( 5,  0,  5), Vector3( 5, 0,  -5)));
    AddQuad( blue,  Quad(Vector3(-5, 10,  5), Vector3(-5, 10, -5), Vector3( 5, 10, -5), Vector3( 5, 10,  5)));
    AddQuad(white,  Quad(Vector3( 5,  0, -5), Vector3( 5, 10, -5), Vector3(-5, 10, -5), Vector3(-5,  0, -5)));
    AddQuad(white,  Quad(Vector3( 5, 10,  5), Vector3( 5,  0,  5), Vector3(-5,  0,  5), Vector3(-5, 10,  5)));
    AddQuad(  red,  Quad(Vector3(-5,  0, -5), Vector3(-5, 10, -5), Vector3(-5, 10,  5), Vector3(-5,  0,  5)));
    AddQuad(green,  Quad(Vector3( 5,  0,  5), Vector3( 5, 10,  5), Vector3( 5, 10, -5), Vector3( 5,  0, -5)));

    AddQuad(white, Quad(Vector3(-2, 5, -1), Vector3(-2, 5,  2), Vector3( 0, 5,  2), Vector3( 0, 5, -1)));
    AddQuad(white, Quad(Vector3(-2, 0,  2), Vector3( 0, 0,  2), Vector3( 0, 5,  2), Vector3(-2, 5,  2)));
    AddQuad(white, Quad(Vector3(-2, 5, -1), Vector3( 0, 5, -1), Vector3( 0, 0, -1), Vector3(-2, 0, -1)));
    AddQuad(white, Quad(Vector3( 0, 5, -1), Vector3( 0, 5,  2), Vector3( 0, 0,  2), Vector3( 0, 0, -1)));
    AddQuad(white, Quad(Vector3(-2, 0, -1), Vector3(-2, 0,  2), Vector3(-2, 5,  2), Vector3(-2, 5, -1)));

    VertexElement elements[] =
    {
        VertexElement(0, VertexElement::kPosition, 3, VertexElement::kFloat),
        VertexElement(sizeof(Vector3), VertexElement::kColour0, 3, VertexElement::kFloat),
    };
    vb = new VertexBuffer(VertexBuffer::kStreaming, elements, 2, sizeof(Vertex), vertices.size());
    SetVertexData(vb, &vertices[0], 0, vertices.size());

    leafPatchCount = allPatches.size();

    ClusterPatches();
    CalculateContributingPatches();

    printf("Calculating contribution debug data...\n");
    contributionDebug = new vector<ContributionDebug>[leafPatchCount];
    for (size_t i=0; i<leafPatchCount; ++i)
    {
        const Patch& patch = *allPatches[i];
        for (vector<Contribution>::const_iterator ic=patch.contributions.begin(), ec=patch.contributions.end(); ic!=ec; ++ic)
            CalculateContributionDebug(ic->patchIdx, contributionDebug[i], *ic);
    }
}


void Render()
{
    glEnable(GL_DEPTH_TEST);
    glDepthMask(GL_TRUE);
    glDepthFunc(GL_LESS);
    glEnable(GL_CULL_FACE);
    glColorMask(GL_TRUE, GL_TRUE, GL_TRUE, GL_TRUE);
    glClearDepth(1.0f);

    glClearColor(0, 0, 0, 0);
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT | GL_STENCIL_BUFFER_BIT);

    BindShader(shader);
    GLint location = glGetUniformLocation(shader->GetHandle(), "worldViewProj");
    glUniformMatrix4fv(location, 1, GL_FALSE, (GLfloat*)&viewProj);

    BindVertexBuffer(vb);
    glDrawArrays(GL_TRIANGLES, 0, vb->GetVertexCount());

    SwapBuffers();
}


void UpdateCamera(double dt)
{
    Vector3 translation = Vector3::kZero;
    if (Keyboard::IsKeyDown(Keyboard::kD))
        translation.x += 1;
    if (Keyboard::IsKeyDown(Keyboard::kA))
        translation.x -=1;
    if (Keyboard::IsKeyDown(Keyboard::kR))
        translation.y += 1;
    if (Keyboard::IsKeyDown(Keyboard::kF))
        translation.y -=1;
    if (Keyboard::IsKeyDown(Keyboard::kS))
        translation.z +=1;
    if (Keyboard::IsKeyDown(Keyboard::kW))
        translation.z -= 1;

    float speed = 5 * (float)dt;
    if (Keyboard::IsKeyDown(Keyboard::kLeftShift))
        speed *= 0.1f;

    Vector3 rotation = Vector3::kZero;
    if (Mouse::IsButtonDown(Mouse::kLeft))
    {
        int dx, dy;
        Mouse::GetMotion(dx, dy);
        rotation.x += 0.01f * dy;
        rotation.y += 0.01f * dx;
    }

    camera.UpdateTransform(speed*translation, rotation);
}


void CalculateDirectIllumination(Vector3 LP)
{
    Vector3 L, P, N;
    Patch** ip = &allPatches[0];
    Patch** ep = ip + leafPatchCount;
    while (ip != ep)
    {
        Patch& patch = **ip++;
        P = patch.position;
        L = LP - P;

        if (IntersectRayQuads(P, LP, &faces[0], faces.size()))
        {
            patch.direct = Vector3::kZero;
        }
        else
        {
            float Ldis = Length(L);
            L /= Ldis;
            N = patch.normal;
            float d = Max(Dot(L, N), 0.0f);
            patch.direct = patch.albedo * (5 * d / (Ldis * Ldis));
        }
        patch.direct += bounceAmount * patch.indirect;
    }
}


void PropagateIllumination()
{
    Vector3 irr;
    Patch** base = &allPatches[0];
    Patch** ip = base + leafPatchCount;
    Patch** ep = base + allPatches.size();
    while (ip != ep)
    {
        Patch& patch = **ip++;
        irr = Vector3::kZero;
        size_t* ic = &patch.children[0];
        size_t* ec = ic + patch.children.size();
        while (ic != ec)
        {
            Patch& child = *allPatches[*ic++];
            irr += child.direct * child.area;
        }
        patch.direct = irr / patch.area;
    }
}


void CalculateIndirectIllumination()
{
    Patch** begin = &allPatches[0];
    Patch** ip = begin;
    Patch** end = begin + allPatches.size();
    if (isIndirectLightingEnabled)
    {
        Vector3 sum;
        while (ip != end)
        {
            Patch& p = **ip++;
            sum = Vector3::kZero;
            for (vector<Contribution>::iterator iq=p.contributions.begin(), eq=p.contributions.end(); iq!=eq; ++iq)
            {
                Contribution& q = *iq;
                sum += begin[q.patchIdx]->direct * q.factor;
            }
            p.indirect = sum * p.albedo;
        }
    }
    else
    {
        while (ip != end)
        {
            Patch& p = **ip++;
            p.indirect = Vector3::kZero;
        }
    }
}


void UpdateVertexBuffer()
{
    size_t j = 0;
    Vector3 colour;
    Patch** ip = &allPatches[0];
    Patch** ep = ip + leafPatchCount;

    if (selectedPatchIdx == kInvalidIndex)
    {
        while (ip != ep)
        {
            const Patch& p = **ip++;
            colour = p.direct + p.indirect;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
        }
    }
    else
    {
        colour = Vector3::kZero;
        while (ip != ep)
        {
            const Patch& p = **ip++;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
        }

        const vector<ContributionDebug>& debug = contributionDebug[selectedPatchIdx];
        for (vector<ContributionDebug>::const_iterator id=debug.begin(), ed=debug.end(); id!=ed; ++id)
        {
            const ContributionDebug& data = *id;
            switch (debugMode)
            {
            case kClusterColour: colour = allPatches[data.contribution.patchIdx]->debugColour; break;
//             case kClusterFormFactor: colour.X = colour.Y = colour.Z = 4 * i.Contribution.Factor / i.Contribution.Patch.Area; break;
//             case kClusterIrradiance: colour = i.Contribution.Patch.Direct / i.Contribution.Patch.Area; break;
//             case kClusterContribution: colour = i.Contribution.Patch.Direct * (100 * i.Contribution.Factor / i.Contribution.Patch.Area); break;
            default: colour = Vector3(1, 1, 0); break;
            }
            j = data.patchIndex * 6;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
            vertices[j++].colour = colour;
        }

        j = selectedPatchIdx * 6;
        colour = Vector3(0, 1, 1);
        vertices[j++].colour = colour;
        vertices[j++].colour = colour;
        vertices[j++].colour = colour;
        vertices[j++].colour = colour;
        vertices[j++].colour = colour;
        vertices[j++].colour = colour;
    }

    SetVertexData(vb, &vertices[0], 0, vertices.size());
}


size_t frame = 0;
double totalElapsed = 0;
void UpdateLighting(Vector3 LP)
{
    double elapsed;
    Timer timer;

    CalculateDirectIllumination(LP);
    PropagateIllumination();
    CalculateIndirectIllumination();

    timer.GetElapsedTime(elapsed);
    totalElapsed += elapsed;
    if (++frame == 30)
    {
        printf("Illumination: %.1fms\n", totalElapsed*1000/30);
        frame = 0;
        totalElapsed = 0;
    }

    UpdateVertexBuffer();
}


bool Update(double dt)
{
    UpdateCamera(dt);

    invProj = InversePerspectiveProjection(60*kDegToRad, kAspectRatio, 0.01f, 100.0f);
    invView = camera.GetTransform();
    invViewProj = invProj * invView;

    proj = PerspectiveProjection(60*kDegToRad, kAspectRatio, 0.01f, 100.0f);
    view = TransformInverseRT(invView);
    viewProj = view * proj;

    if (Keyboard::IsKeyDown(Keyboard::kC))
        selectedPatchIdx = kInvalidIndex;
    if (Keyboard::WasKeyPressed(Keyboard::kSpace))
        isPaused = !isPaused;

    if (Keyboard::IsKeyDown(Keyboard::k0))
        isIndirectLightingEnabled = false;
    if (Keyboard::IsKeyDown(Keyboard::k1))
    {
        isIndirectLightingEnabled = true;
        bounceAmount = 0;
    }
    if (Keyboard::IsKeyDown(Keyboard::k2))
    {
        isIndirectLightingEnabled = true;
        bounceAmount = 0.5f;
    }
    if (Keyboard::IsKeyDown(Keyboard::k3))
    {
        isIndirectLightingEnabled = true;
        bounceAmount = 0.75f;
    }
    if (Keyboard::IsKeyDown(Keyboard::k4))
    {
        isIndirectLightingEnabled = true;
        bounceAmount = 1;
    }

    if (Mouse::IsButtonDown(Mouse::kRight))
    {
        int x, y;
        Mouse::GetPosition(x, y);
        Vector4 p = Vector4((2.0f*x)/(float)kScreenWidth-1, (2.0f*y)/(float)kScreenHeight-1, 0, 1);
        Vector4 q = Vector4((2.0f*x)/(float)kScreenWidth-1, 1-(2.0f*y)/(float)kScreenHeight, 1, 1);
        p = MulVectorMatrix(p, invViewProj);
        q = MulVectorMatrix(q, invViewProj);
        p /= p.w;
        q /= q.w;
        selectedPatchIdx = IntersectRayQuadsClosest(p.AsVector3(), q.AsVector3(), &patchQuads[0], patchQuads.size());
    }

    if (!isPaused)
    {
        elapsedTime += dt;
        lightPos.x = 4.5f * (float)cos(0.3 * elapsedTime);
        lightPos.y = 5.0f + 4.5f * (float)sin(0.4 * elapsedTime);
        lightPos.z = 4.5f * (float)sin(0.5 * elapsedTime);
    }
    UpdateLighting(lightPos);
    
    bool isDone = Keyboard::IsKeyDown(Keyboard::kEscape);
    return isDone;
}


bool Startup(int, char*[])
{
	CreateDevice("GI", kScreenWidth, kScreenHeight, kCreateDepthBuffer, 8);
    InitScene();
	return true;
}


void Shutdown()
{
    delete vb;
    delete shader;
}
